from flask import Flask, Response, jsonify
from flask_cors import CORS
import cv2
import time
import threading
import base64
import numpy as np
from pygrabber.dshow_graph import FilterGraph
import json

app = Flask(__name__)
CORS(app)  # 允许跨域请求

class VideoStreamServer:
    def __init__(self):
        self.cap = None
        self.is_streaming = False
        self.current_frame = None
        self.fps = 0
        self.frame_count = 0
        self.start_time = time.time()
        
    def initialize_camera(self, camera_name="USB Camera"):
        """初始化摄像头"""
        try:
            graph = FilterGraph()
            devices = graph.get_input_devices()
            
            if camera_name in devices:
                index = devices.index(camera_name)
                self.cap = cv2.VideoCapture(index, cv2.CAP_DSHOW)
                
                # 设置摄像头参数
                self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)  # 降低分辨率以提高性能
                self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
                self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M','J','P','G'))
                self.cap.set(cv2.CAP_PROP_FPS, 30)
                
                if self.cap.isOpened():
                    print(f"摄像头初始化成功: {camera_name}")
                    return True
                else:
                    print("无法打开摄像头")
                    return False
            else:
                print(f"未找到摄像头: {camera_name}")
                print(f"可用设备: {devices}")
                return False
        except Exception as e:
            print(f"摄像头初始化失败: {e}")
            return False
    
    def start_streaming(self):
        """开始视频流"""
        if not self.cap or not self.cap.isOpened():
            if not self.initialize_camera():
                return False
        
        self.is_streaming = True
        self.start_time = time.time()
        
        def stream_loop():
            while self.is_streaming:
                ret, frame = self.cap.read()
                if ret:
                    self.current_frame = frame
                    self.frame_count += 1
                    
                    # 计算 FPS
                    elapsed_time = time.time() - self.start_time
                    if elapsed_time > 0:
                        self.fps = self.frame_count / elapsed_time
                else:
                    print("无法读取帧")
                    break
                time.sleep(0.01)  # 10ms 延迟
        
        threading.Thread(target=stream_loop, daemon=True).start()
        return True
    
    def stop_streaming(self):
        """停止视频流"""
        self.is_streaming = False
        if self.cap:
            self.cap.release()
    
    def get_frame_base64(self):
        """获取当前帧的 base64 编码"""
        if self.current_frame is not None:
            # 压缩图像以提高传输效率
            encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 80]
            _, buffer = cv2.imencode('.jpg', self.current_frame, encode_param)
            return base64.b64encode(buffer).decode('utf-8')
        return None
    
    def get_frame_info(self):
        """获取帧信息"""
        if self.current_frame is not None:
            height, width = self.current_frame.shape[:2]
            return {
                'width': width,
                'height': height,
                'fps': round(self.fps, 2),
                'frame_count': self.frame_count,
                'timestamp': time.time()
            }
        return None

# 全局视频流服务器实例
video_server = VideoStreamServer()

@app.route('/')
def index():
    return jsonify({
        'status': 'Video Stream Server Running',
        'endpoints': {
            '/start': 'Start video streaming',
            '/stop': 'Stop video streaming',
            '/frame': 'Get current frame as base64',
            '/info': 'Get frame information',
            '/status': 'Get server status'
        }
    })

@app.route('/start')
def start_stream():
    """启动视频流"""
    if video_server.start_streaming():
        return jsonify({'status': 'success', 'message': '视频流已启动'})
    else:
        return jsonify({'status': 'error', 'message': '无法启动视频流'}), 500

@app.route('/stop')
def stop_stream():
    """停止视频流"""
    video_server.stop_streaming()
    return jsonify({'status': 'success', 'message': '视频流已停止'})

@app.route('/frame')
def get_frame():
    """获取当前帧"""
    frame_data = video_server.get_frame_base64()
    if frame_data:
        return jsonify({
            'status': 'success',
            'frame': frame_data,
            'timestamp': time.time()
        })
    else:
        return jsonify({'status': 'error', 'message': '无可用帧'}), 404

@app.route('/info')
def get_info():
    """获取帧信息"""
    info = video_server.get_frame_info()
    if info:
        return jsonify({'status': 'success', 'info': info})
    else:
        return jsonify({'status': 'error', 'message': '无可用信息'}), 404

@app.route('/status')
def get_status():
    """获取服务器状态"""
    return jsonify({
        'status': 'success',
        'is_streaming': video_server.is_streaming,
        'camera_opened': video_server.cap.isOpened() if video_server.cap else False,
        'fps': round(video_server.fps, 2),
        'frame_count': video_server.frame_count
    })

if __name__ == '__main__':
    print("启动视频流服务器...")
    print("访问 http://localhost:5000 查看API文档")
    app.run(host='0.0.0.0', port=5000, debug=False, threaded=True) 